#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
EP加密验证 - 完整捕获RSA密钥和加密流程
"""

import frida
import sys
import time
import io

# Fix encoding for Windows console
sys.stdout = io.TextIOWrapper(sys.stdout.buffer, encoding='utf-8', errors='replace')

print("\n" + "="*80)
print("🔑 EP加密验证 - RSA密钥和加密流程完整追踪")
print("="*80 + "\n")

# ============================================================================
# 1. 连接设备
# ============================================================================
print("[1] 连接设备...")
try:
    device = frida.get_usb_device()
    print(f"    ✓ 已连接: {device}")
except Exception as e:
    print(f"    ✗ 连接失败: {e}")
    print("\n请确保:")
    print("  1. 手机已通过USB连接")
    print("  2. 已开启USB调试")
    print("  3. adb devices能看到设备")
    print("  4. frida-server正在运行")
    sys.exit(1)

# ============================================================================
# 2. 选择应用
# ============================================================================
print("\n[2] 选择应用...")
print("\n请选择启动方式:")
print("  1. 启动新进程 (推荐)")
print("  2. 附加到运行中的进程")

choice = input("\n输入选项 [1/2]: ").strip()

if choice == "1":
    # 启动新进程
    package_name = input("\n输入应用包名 (如 com.mcdonalds.gma.cn): ").strip()
    
    if not package_name:
        print("✗ 包名不能为空")
        sys.exit(1)
    
    try:
        print(f"\n启动应用: {package_name}...")
        pid = device.spawn([package_name])
        print(f"    ✓ 已启动进程 PID: {pid}")
        
        session = device.attach(pid)
        print(f"    ✓ 已附加到进程")
        
        should_resume = True
        
    except Exception as e:
        print(f"    ✗ 启动失败: {e}")
        sys.exit(1)
        
elif choice == "2":
    # 附加到运行中的进程
    print("\n正在获取进程列表...")
    
    try:
        processes = device.enumerate_applications()
        apps = [app for app in processes if not app.identifier.startswith('com.android')]
        
        print("\n已安装的应用:")
        for idx, app in enumerate(apps[:20], 1):
            print(f"  {idx}. {app.name} ({app.identifier})")
        
        app_idx = int(input("\n选择应用编号: ").strip()) - 1
        
        if app_idx < 0 or app_idx >= len(apps):
            print("✗ 无效的编号")
            sys.exit(1)
        
        target_app = apps[app_idx]
        package_name = target_app.identifier
        
        print(f"\n附加到: {target_app.name} ({package_name})...")
        session = device.attach(package_name)
        print(f"    ✓ 已附加")
        
        should_resume = False
        
    except Exception as e:
        print(f"    ✗ 附加失败: {e}")
        sys.exit(1)
else:
    print("✗ 无效的选项")
    sys.exit(1)

# ============================================================================
# 3. 加载Hook脚本
# ============================================================================
print("\n[3] 加载Hook脚本...")

script_file = 'verify_rsa_key.js'

try:
    with open(script_file, 'r', encoding='utf-8') as f:
        script_code = f.read()
except FileNotFoundError:
    print(f"    ✗ 找不到脚本文件: {script_file}")
    print(f"    请确保 {script_file} 在当前目录下")
    sys.exit(1)
except Exception as e:
    print(f"    ✗ 读取脚本失败: {e}")
    sys.exit(1)

script = session.create_script(script_code)
print(f"    ✓ 脚本已创建")

# 消息处理器
def on_message(message, data):
    if message['type'] == 'send':
        payload = message['payload']
        print(payload)
    elif message['type'] == 'error':
        print(f"\n[FRIDA ERROR] {message.get('description', message)}")
        if 'stack' in message:
            print(f"Stack:\n{message['stack']}")
    else:
        print(message)

script.on('message', on_message)

try:
    script.load()
    print(f"    ✓ 已加载: {script_file}")
except Exception as e:
    print(f"    ✗ 加载失败: {e}")
    sys.exit(1)

# 如果是spawn模式，恢复进程
if should_resume:
    device.resume(pid)
    print(f"    ✓ 已恢复进程运行")

# ============================================================================
# 4. 等待Hook初始化
# ============================================================================
print("\n[4] 等待Hook初始化...")
time.sleep(3)
print("    ✓ Hook已就绪")

# ============================================================================
# 5. 等待用户操作
# ============================================================================
print("\n" + "="*80)
print("✅ Hook已激活！")
print("="*80)
print("\n现在请在应用中执行以下操作:")
print("  1. 登录账号")
print("  2. 提交表单")
print("  3. 或任何会触发数美SDK的操作")
print("\n观察下面的输出，会显示捕获的密钥信息")
print("\n按 Ctrl+C 停止...")
print("="*80 + "\n")

# ============================================================================
# 6. 保持运行
# ============================================================================
try:
    sys.stdin.read()
except KeyboardInterrupt:
    print("\n\n[+] 用户中断")

# ============================================================================
# 7. 清理
# ============================================================================
print("\n[+] 停止Hook...")
try:
    script.unload()
    session.detach()
    print("[OK] 已退出\n")
except:
    pass

